"""
@author: chenzhenhua
@project: jf_fashion
@file: mnist.py
@time: 2021/8/3 0003 8:50
@desc:
"""

import numpy as np
from tensorflow.python.keras.utils import np_utils  # keras中的numpy工具包


def load_data(path='mnist.npz'):
    """Loads the MNIST dataset.

    # Arguments
        path: path where to cache the dataset locally
            (relative to ~/.keras/datasets).

    # Returns
        Tuple of Numpy arrays: `(x_train, y_train), (x_test, y_test)`.
    """
    f = np.load(path)
    x_train, y_train = f['x_train'], f['y_train']
    x_test, y_test = f['x_test'], f['y_test']
    return (x_train, y_train), (x_test, y_test)


def read_mnist(path):
    return load_data(path)


def preprocessing(x1, y1, x2, y2, mode=0):
    if mode==0:
        x_train = x1.reshape(x1.shape[0], -1) / 255.0
        x_test = x2.reshape(x2.shape[0], -1) / 255.0
    elif mode==1:
        x_train = x1 / 255.0
        x_test = x2 / 255.0
    else:
        x_train = x1.reshape(x1.shape[0], 28, 28, 1)
        x_test = x2.reshape(x2.shape[0], 28, 28, 1)
    # 转换为one hot格式
    y_train = np_utils.to_categorical(y1, num_classes=10)
    y_test = np_utils.to_categorical(y2, num_classes=10)
    return x_train, y_train, x_test, y_test
